/*
 * Environment:
 * Microsoft Visual Studio 2008 Version 9.0.21022.8 RTM
 * Microsoft (R) 32-bit C/C++ Optimizing Compiler Version 15.00.21022.08 for 80x86
 *
 * Authors: Sergey Zamyslov, Sergey Moshkov
 */

#include <math.h>
#include <stdio.h>
#include <stdlib.h>
#include <time.h>

#include "tbb/task_scheduler_init.h"
#include "tbb/blocked_range.h"
#include "tbb/parallel_for.h"
#include "tbb/tick_count.h"

using namespace std;
using namespace tbb;

#define DEBUG        false
#define THREAD_COUNT 4
#define GRAINSIZE    1

/*** Help function ***/

double* clone(double* from, int size)
{
    int i;
    double* clone = (double*)malloc(sizeof(double) * size);

    for(i = 0; i < size; ++i)
    {
        clone[i] = from[i];
    }

    return clone;
}

/*** ---end--- ***/

/*** Function to interpolate ***/

double function(double x) {
    return cos(x) / x;
}

/*** ---end--- ***/

/*** Output functions ***/

void printExtra(double* data, int n, int extra)
{
    char* format      = "%11.4f";
    char* extraFormat = "%24.18f\n";

    int i;

    if(n > 20)
    {
        return;
    }

    if(extra)
    {
        format = extraFormat;
    }

    for (i = 0; i < n; i++)
    {
        printf(format, data[i]);
    }
}

void print(double* data, int n)
{
    printExtra(data, n, 0);
}

/*** ---end--- ***/

/*** TDMA Serial ***/

double solveTDM(int equationCount,
                double* subDiagonal, double* mainDiagonal, double* supDiagonal,
                double* b,
                double* x)
{
    
	// Initialization
	clock_t start, finish;

	start = clock();

    int     i;
    size_t	sizeInBytes      = sizeof(double) * equationCount;
    double* mainDiagonalWork = (double*)malloc(sizeInBytes),
          * bWork            = (double*)malloc(sizeInBytes);

	mainDiagonalWork[0] = mainDiagonal[0];
	bWork[0]            = b[0];

    // Forward
    for (i = 1; i < equationCount; i++)
    {
        double m = subDiagonal[i-1] / mainDiagonalWork[i-1];

        mainDiagonalWork[i] = mainDiagonal[i] - m * supDiagonal[i-1];
        bWork[i]            = b[i] - m * bWork[i-1];
	}

    // Backward
    x[equationCount - 1] = bWork[equationCount - 1] / mainDiagonalWork[equationCount - 1];	
    for (i = equationCount - 2; i >= 0; i--)
    {
        x[i] = (bWork[i] - supDiagonal[i] * x[i + 1]) / mainDiagonalWork[i];
    }

    free(mainDiagonalWork);
    free(bWork);

	finish = clock();

	return (((double)finish - (double)start) / CLOCKS_PER_SEC);
}

/*** ---end--- ***/

/**************************************/
/*** TDMA Block TBB implementation  ***/
/*** (one thread handles one block) ***/
/**************************************/

/*** Forward step ***/

void doForward(double *subDiagonal, double *mainDiagonal, double *supDiagonal, double *b, int blockRowCount, int blockIndex)
{
	if(DEBUG)
    {
        printf("Block #%d is here.\n", blockIndex);
    }

    int blockStartRow = blockIndex * blockRowCount;
    int blockEndRow   = blockStartRow + blockRowCount - 1;

    for(int blockRowIndex = blockStartRow; blockRowIndex < blockEndRow; blockRowIndex++)
    {
        if(DEBUG)
        {
            printf("Processing row %d.\n", blockRowIndex);
        }

        double factor = subDiagonal[blockRowIndex] / mainDiagonal[blockRowIndex];

        if(blockIndex > 0)
        {
            // Non-diagonal elements.
            subDiagonal[blockRowIndex] = -subDiagonal[blockRowIndex - 1] * factor;
        }
        else
        {
            subDiagonal[blockRowIndex] = 0.0;
        }

        mainDiagonal[blockRowIndex + 1] -= supDiagonal[blockRowIndex] * factor;
        b[blockRowIndex + 1] -= b[blockRowIndex] * factor;
    }

    if(DEBUG)
    {
        if(blockIndex > 0)
        {
            printf("Left non-diagonal elements:\n");
            print(&subDiagonal[blockStartRow], blockRowCount - 1);
        }
        printf("\n\n");
    }
}

class BlockForward
{
private:
	double *subDiagonal, *mainDiagonal, *supDiagonal, *b;
	int const blockRowCount;

public:
	BlockForward(double *_subDiagonal, double *_mainDiagonal, double *_supDiagonal, double *_b, int _blockRowCount) :
		subDiagonal(_subDiagonal), mainDiagonal(_mainDiagonal), supDiagonal(_supDiagonal), b(_b), blockRowCount(_blockRowCount)
	{}

	void operator() (const blocked_range<int>& r) const
	{
		for (int i = r.begin(); i != r.end(); i++)
		{
			doForward(subDiagonal, mainDiagonal, supDiagonal, b, blockRowCount, i);
		}
	}
};

/*** ---end--- ***/

/*** Backward step ***/

void doBackward(double *subDiagonal, double *mainDiagonal, double *supDiagonal, double *b, int blockRowCount, int blockIndex)
{
	if(DEBUG)
    {
        printf("Block #%d is here.\n", blockIndex);
    }

    int blockStartRow = blockIndex * blockRowCount;    
	int blockEndRow   = blockStartRow + blockRowCount - 1;
		blockStartRow = max(blockStartRow, 1);

    for (int blockRowIndex = blockEndRow - 1; blockRowIndex >= blockStartRow; blockRowIndex--)
    {
        if(DEBUG)
        {
            printf("Processing row %d.\n", blockRowIndex);
        }

        double factor = supDiagonal[blockRowIndex - 1] / mainDiagonal[blockRowIndex];

        if(blockIndex > 0)
        {
            if(blockRowIndex > blockStartRow)
            {
                subDiagonal[blockRowIndex - 2] -= subDiagonal[blockRowIndex - 1] * factor;
            }
            else
            {
                // If we edit element out of block, we take element from main diagonal.
                mainDiagonal[blockRowIndex - 1] -= subDiagonal[blockRowIndex - 1] * factor;
            }
        }

        // Non-diagonal elements.
        supDiagonal[blockRowIndex - 1] = -supDiagonal[blockRowIndex] * factor;

        b[blockRowIndex - 1] -= b[blockRowIndex] * factor;
    }

    if(DEBUG)
    {
        printf("\n\n");
    }
}

class BlockBackward
{
private:
	double *subDiagonal, *mainDiagonal, *supDiagonal, *b;
	int const blockRowCount;

public:
	BlockBackward(double *_subDiagonal, double *_mainDiagonal, double *_supDiagonal, double *_b, int _blockRowCount) :
		subDiagonal(_subDiagonal), mainDiagonal(_mainDiagonal), supDiagonal(_supDiagonal), b(_b), blockRowCount(_blockRowCount)
	{}

	void operator() (const blocked_range<int>& r) const
	{
		for (int i = r.begin(); i != r.end(); i++)
		{
			doBackward(subDiagonal, mainDiagonal, supDiagonal, b, blockRowCount, i);
		}
	}
};

/*** ---end--- ***/

/*** Make an auxilary matrix step ***/

void doAuxilaryMatrix(double *subDiagonal, double *mainDiagonal, double *supDiagonal, double *b, int blockRowCount, int blockIndex, double **auxiliaryMatrix, int blockCount)
{
	int blockStartRow = blockIndex * blockRowCount;
    int blockEndRow   = blockStartRow + blockRowCount - 1;

    if(blockIndex > 0)
    {
        auxiliaryMatrix[0][blockIndex - 1] = subDiagonal[blockEndRow - 1];
    }

    if(blockIndex < blockCount - 1)
    {
        auxiliaryMatrix[2][blockIndex] = supDiagonal[blockEndRow];
    }

    auxiliaryMatrix[1][blockIndex] = mainDiagonal[blockEndRow];
    auxiliaryMatrix[3][blockIndex] = b[blockEndRow];
}

class BlockAuxilaryMatrix
{
private:
	double *subDiagonal, *mainDiagonal, *supDiagonal, *b, **auxilaryMatrix;
	int const blockRowCount;
	int const blockCount;

public:
	BlockAuxilaryMatrix(double *_subDiagonal, double *_mainDiagonal, double *_supDiagonal, double *_b, int _blockRowCount, double **_auxilaryMatrix, int _blockCount) :
		subDiagonal(_subDiagonal), mainDiagonal(_mainDiagonal), supDiagonal(_supDiagonal), b(_b), blockRowCount(_blockRowCount), auxilaryMatrix(_auxilaryMatrix), blockCount(_blockCount)
	{}

	void operator() (const blocked_range<int>& r) const
	{
		for (int i = r.begin(); i != r.end(); i++)
		{
			doAuxilaryMatrix(subDiagonal, mainDiagonal, supDiagonal, b, blockRowCount, i, auxilaryMatrix, blockCount);
		}
	}
};

/*** ---end--- ***/

/*** Set X step ***/

void doX(double *auxiliaryX, double *x, int blockRowCount, int blockIndex)
{
    int blockStartRow = blockIndex * blockRowCount;
    int blockEndRow   = blockStartRow + blockRowCount - 1;

    if(DEBUG)
    {
        for(int blockRowIndex = blockStartRow; blockRowIndex < blockEndRow; blockRowIndex++)
        {
            x[blockRowIndex] = 0.0;
        }
    }

    x[blockEndRow] = auxiliaryX[blockIndex];
}

class BlockX
{
private:
	double *auxiliaryX, *x;
	int const blockRowCount;

public:
	BlockX(double *_auxiliaryX, double *_x, int _blockRowCount) :
		auxiliaryX(_auxiliaryX), x(_x), blockRowCount(_blockRowCount)
	{}

	void operator() (const blocked_range<int>& r) const
	{
		for (int i = r.begin(); i != r.end(); i++)
		{
			doX(auxiliaryX, x, blockRowCount, i);
		}
	}
};

/*** ---end--- ***/

/*** Calculate X step ***/

void doCalculate(double *subDiagonal, double *mainDiagonal, double *supDiagonal, double *b, double *x, int blockRowCount, int blockIndex)
{
	int blockStartRow = blockIndex * blockRowCount;
    int blockEndRow   = blockStartRow + blockRowCount - 1;

    for(int blockRowIndex = blockStartRow; blockRowIndex < blockEndRow; blockRowIndex++)
    {
        x[blockRowIndex] = b[blockRowIndex];

        if(blockIndex > 0)
        {
            x[blockRowIndex] -= subDiagonal[blockRowIndex - 1] * x[blockStartRow - 1];
        }

        x[blockRowIndex] -= supDiagonal[blockRowIndex] * x[blockEndRow];

        x[blockRowIndex] /= mainDiagonal[blockRowIndex];
    }
}

class BlockCalculate
{
private:
	double *subDiagonal, *mainDiagonal, *supDiagonal, *b, *x;
	int const blockRowCount;

public:
	BlockCalculate(double *_subDiagonal, double *_mainDiagonal, double *_supDiagonal, double *_b, double *_x, int _blockRowCount) :
		subDiagonal(_subDiagonal), mainDiagonal(_mainDiagonal), supDiagonal(_supDiagonal), b(_b), x(_x), blockRowCount(_blockRowCount)
	{}

	void operator() (const blocked_range<int>& r) const
	{
		for (int i = r.begin(); i != r.end(); i++)
		{
			doCalculate(subDiagonal, mainDiagonal, supDiagonal, b, x, blockRowCount, i);
		}
	}
};

/*** ---end--- ***/

/*** solveTDM_Block ***/

double solveTDM_Block(int equationCount,
					  double* _subDiagonal, double* _mainDiagonal, double* _supDiagonal,
					  double* _b,
					  int blockCount,
					  double* x)
{
    if (equationCount % blockCount != 0)
	{
        printf("%d equations cannot be divided in %d blocks.\n", equationCount, blockCount);
        exit(-1);
    }
	
	size_t dataElementSize = sizeof(double);

    int blockRowCount = equationCount / blockCount;

    int blockIndex;
    int blockStartRow, blockEndRow;
    int blockRowIndex;

    double  factor;
    double* auxiliaryMatrix[5];

	tick_count start, finish;

    // Copy all data.
    double *subDiagonal  = clone(_subDiagonal, equationCount - 1);
    double *mainDiagonal = clone(_mainDiagonal, equationCount);
    double *supDiagonal  = clone(_supDiagonal, equationCount - 1);
    double *b			 = clone(_b, equationCount);

	start = tick_count::now();

    /*** Step one ***/
    if(DEBUG)
    {
        printf("**********************************\n");
        printf("Step one. Forward step. \n\n");
    }

    parallel_for(blocked_range<int>(0, blockCount, GRAINSIZE), BlockForward(subDiagonal, mainDiagonal, supDiagonal, b, blockRowCount));
	/*** end of step one ***/

    // Synchronization there is automatic parallel_for synchronization

    /*** Step two ***/
    if(DEBUG)
    {
        printf("**********************************\n");
        printf("Step two. Backward step. \n\n");
    }

    parallel_for(blocked_range<int>(0, blockCount, GRAINSIZE), BlockBackward(subDiagonal, mainDiagonal, supDiagonal, b, blockRowCount));
	/*** end of step two ***/

    // Synchronization there is automatic parallel_for synchronization

	/*** Step three ***/
	if(DEBUG)
    {
        printf("**********************************\n");
        printf("Step three. Form auxilary matrix. \n\n");
    }
	
	auxiliaryMatrix[0] = (double*)malloc(dataElementSize * (blockCount - 1));
	auxiliaryMatrix[1] = (double*)malloc(dataElementSize * blockCount);
	auxiliaryMatrix[2] = (double*)malloc(dataElementSize * (blockCount - 1));
	auxiliaryMatrix[3] = (double*)malloc(dataElementSize * blockCount);
	auxiliaryMatrix[4] = (double*)malloc(dataElementSize * blockCount);

	parallel_for(blocked_range<int>(0, blockCount, GRAINSIZE), BlockAuxilaryMatrix(subDiagonal, mainDiagonal, supDiagonal, b, blockRowCount, auxiliaryMatrix, blockCount));
	/*** end of step three ***/

	// Synchronization there is automatic parallel_for synchronization

	/*** Step four ***/
	if(DEBUG)
    {
        printf("**********************************\n");
        printf("Step four. Solve auxilary matrix. \n\n");
    }

    solveTDM(blockCount,
			 auxiliaryMatrix[0], auxiliaryMatrix[1], auxiliaryMatrix[2],
			 auxiliaryMatrix[3],
			 auxiliaryMatrix[4]);
	/*** end of step four ***/

    /*** Step five ***/
	if(DEBUG)
    {
        printf("**********************************\n");
        printf("Step five. Form X. \n\n");
    }

    parallel_for(blocked_range<int>(0, blockCount, GRAINSIZE), BlockX(auxiliaryMatrix[4], x, blockRowCount));
	
    if(DEBUG)
    {
        printf("X:\n");
        print(x, equationCount);
        printf("\n\n\n");
    }
	/*** end of step five ***/
	
	// Synchronization there is automatic parallel_for synchronization
    
	/*** Step six ***/
	if(DEBUG)
    {
        printf("**********************************\n");
        printf("Step six. Calculate X. \n\n");
    }

	parallel_for(blocked_range<int>(0, blockCount, GRAINSIZE), BlockCalculate(subDiagonal, mainDiagonal, supDiagonal, b, x, blockRowCount));
	/*** end of step six ***/

	// Synchronization there is automatic parallel_for synchronization

	free(auxiliaryMatrix[0]);
    free(auxiliaryMatrix[1]);
    free(auxiliaryMatrix[2]);
    free(auxiliaryMatrix[3]);
    free(auxiliaryMatrix[4]);

	finish = tick_count::now();

	free(subDiagonal);
	free(mainDiagonal);
	free(supDiagonal);
	free(b);

    return (finish - start).seconds();
}

/*** ---end--- ***/

/********************************************/
/*** end of TDMA Block TBB implementation ***/
/********************************************/

/*** Additional help functions ***/

void initializeKnotsAndNodes(double a, double b,
                             double* knots, int knotCount,
                             double* nodesX, double* nodesY)
{
    int    i;
    double step = (b - a) / (knotCount - 1);

    double d = step * 0.9;
    double dOver2 = d / 2.0;

    knots[0]             = a;
    knots[knotCount - 1] = b;

    nodesX[0] = knots[0];
    nodesY[0] = function(nodesX[0]);

    srand((int)time(NULL));

    for (i = 1; i < knotCount - 1; i++)
    {
        knots[i] = i * step;
        knots[i] += (((double)rand() / RAND_MAX) * d) - dOver2;

        nodesX[i] = (knots[i] + knots[i - 1]) / 2;
        nodesY[i] = function(nodesX[i]);
    }

    i = knotCount - 1;
    nodesX[i] = (knots[i] + knots[i - 1]) / 2;
    nodesY[i] = function(nodesX[i]);

    nodesX[knotCount] = knots[knotCount - 1];
    nodesY[knotCount] = function(nodesX[knotCount]);
}

// Simple random initializer
void initializeData(int equationCount,
                    double* subDiagonal, double* mainDiagonal, double* supDiagonal,
                    double* b,
                    int quaterMax)
{
    int    i;
    srand((int)time(NULL));

    for (i = 0; i < equationCount - 1; i++)
    {
        subDiagonal[i]  = ((double)rand() / RAND_MAX) * quaterMax;
        supDiagonal[i]  = ((double)rand() / RAND_MAX) * quaterMax;
        mainDiagonal[i] = ((double)rand() / RAND_MAX) * 2 * quaterMax + (quaterMax * 2);
        b[i]            = ((double)rand() / RAND_MAX) * 4 * quaterMax;
    }
    mainDiagonal[i]     = ((double)rand() / RAND_MAX) * 2 * quaterMax + (quaterMax * 2);
    b[i]                = ((double)rand() / RAND_MAX) * 4 * quaterMax;
}

void calculateResiduals(int equationCount,
                        double* subDiagonal, double* mainDiagonal, double* supDiagonal,
                        double* b,
                        double* x,
                        double* residuals)
{
    int i;

    // First row.
    i = 0;
    residuals[i] = mainDiagonal[i] * x[i];
    residuals[i] += supDiagonal[i] * x[i + 1];
    residuals[i] -= b[i];

    for(i = 1; i < equationCount - 1; ++i)
    {
        residuals[i] = mainDiagonal[i] * x[i];

        residuals[i] += subDiagonal[i - 1] * x[i - 1];
        residuals[i] += supDiagonal[i] * x[i + 1];

        residuals[i] -= b[i];
    }

    // Last row.
    i = equationCount - 1;
    residuals[i] = mainDiagonal[i] * x[i];
    residuals[i] += subDiagonal[i - 1] * x[i - 1];
    residuals[i] -= b[i];
}

double getSplinePointY(int knotCount, double* knots,
                       double* nodesX, double* nodesY,
                       double* coeffs,
                       double x)
{
    int i;
    double c;

    for (i = 1; i < knotCount; ++i)
    {
        if (x <= knots[i])
        {
            break;
        }
    }

    if (i < knotCount)
    {
        c = x - nodesX[i];

        return nodesY[i] +
            0.5 * (coeffs[i] + coeffs[i - 1]) * c + 
            (coeffs[i] - coeffs[i - 1]) * c * c / (2 * (knots[i] - knots[i - 1]));
    }

    return 0.0;
}

void calculateSplineCoeffs(double* knots, int knotCount, double* nodesY,
                           int blockCount, double* coeffs)
{
    double* subDiagonal  = (double*)malloc((knotCount - 1) * sizeof(double)),
          * supDiagonal  = (double*)malloc((knotCount - 1) * sizeof(double)),
          * mainDiagonal = (double*)malloc(knotCount * sizeof(double)),
          * b            = (double*)malloc(knotCount * sizeof(double));

    int i;
    double time;

    for (i = 1; i < knotCount - 1; i++)
    {
        subDiagonal[i - 1] = knots[i] - knots[i - 1];
        supDiagonal[i - 1] = knots[i] - knots[i - 1];
        mainDiagonal[i]    = 3 * (knots[i + 1] - knots[i - 1]);
        b[i - 1]           = 8 * (nodesY[i] - nodesY[i - 1]);
    }

    mainDiagonal[0]             = 3 * (knots[1] - knots[0]);
    mainDiagonal[knotCount - 1] = 3 * (knots[knotCount - 1] - knots[knotCount - 2]);
    subDiagonal[knotCount - 2]  = knots[knotCount - 1] - knots[knotCount - 2];
    supDiagonal[knotCount - 2]  = knots[knotCount - 1] - knots[knotCount - 2];
    b[knotCount - 2]            = 8 * (nodesY[knotCount - 1] - nodesY[knotCount - 2]);
    b[knotCount - 1]            = 8 * (nodesY[knotCount] - nodesY[knotCount - 1]);

    time = solveTDM_Block(knotCount, subDiagonal, mainDiagonal, supDiagonal, b, blockCount, coeffs);
    printf("Time: %f\n\n\n", time);

    free(subDiagonal);
    free(supDiagonal);
    free(mainDiagonal);
    free(b);
}

/*** ---end--- ***/

int main(int argc, char* argv[])
{
	task_scheduler_init init;

	init.initialize(THREAD_COUNT);

    double  intervalStart = 0.5;
    double  intervalEnd   = 1000000000.0;
    int     knotCount     = 10000000;
    double* knots         = (double*)malloc(knotCount * sizeof(double));
    double* coeffs        = (double*)malloc(knotCount * sizeof(double));
    double* nodesX        = (double*)malloc((knotCount + 1) * sizeof(double));
    double* nodesY        = (double*)malloc((knotCount + 1) * sizeof(double));

	bool   output         = false;
	bool   outputNodes    = false;
	char*  outputFileName = "chart.csv";
	double chartDx		  = (intervalEnd - intervalStart) / 1000.0;
    
	initializeKnotsAndNodes(intervalStart, intervalEnd, knots, knotCount, nodesX, nodesY);
    calculateSplineCoeffs(knots, knotCount, nodesY, THREAD_COUNT, coeffs);

	// Output.
    if(output)
    {
        FILE* outputFile = fopen(outputFileName, "w+");
	
		fprintf(outputFile, "X;Y;splineY;nodesY\n");
        double chartX;
        for(chartX = intervalStart; chartX < intervalEnd; chartX += chartDx) {
            fprintf(outputFile, "%f;%f;%f;NA\n",
					chartX,
					function(chartX),
					getSplinePointY(knotCount, knots, nodesX, nodesY, coeffs, chartX));
        }

        chartX = intervalEnd;
        fprintf(outputFile, "%f;%f;%f;NA\n", chartX,
				function(chartX),
				getSplinePointY(knotCount, knots, nodesX, nodesY, coeffs, chartX));

        if(outputNodes)
        {
            for(int i = 0; i <= knotCount; ++i) {
                fprintf(outputFile, "%f;NA;NA;%f\n", nodesX[i], nodesY[i]);
            }
        }

        fclose(outputFile);
    }
	
	free(knots);
    free(coeffs);
    free(nodesX);
    free(nodesY);

	init.terminate();

    return 0;
}
